// SPDX-FileCopyrightText: Copyright 2025-2025 深圳市同心圆网络有限公司
// SPDX-License-Identifier: GPL-3.0-only

package trace_dao

import (
	"bytes"
	"encoding/binary"
	"fmt"
	"strings"
	"time"

	"gitcode.com/opendragonfly/df_proto_gen_go.git/trace_api"
	"github.com/dgraph-io/badger/v4"
)

const (
	_START_TIME_INDEX_NAME_SPACE     = "start"
	_START_TIME_ALL_INDEX_NAME_SPACE = "startall"
)

type StartTimeKey struct {
	ServiceName    string //服务名
	ServiceVersion string //服务版本
	RootSpanName   string //根Span名称
	StartTime      int64  //开始时间
}

func (key *StartTimeKey) ToBytes(withTime bool) ([]byte, error) {
	buf := &bytes.Buffer{}

	ns := _START_TIME_INDEX_NAME_SPACE
	if key.RootSpanName == "" {
		ns = _START_TIME_ALL_INDEX_NAME_SPACE
	}
	name := strings.Join([]string{ns, key.ServiceName, key.ServiceVersion, key.RootSpanName}, "\t")
	_, err := buf.Write([]byte(name))
	if err != nil {
		return nil, err
	}
	if withTime {
		timeBytes := make([]byte, 8)
		binary.BigEndian.PutUint64(timeBytes, uint64(key.StartTime))
		_, err = buf.Write(timeBytes)
		if err != nil {
			return nil, err
		}
	}
	return buf.Bytes(), nil
}

func ParseStartTimeKey(buf []byte, withTime bool) (*StartTimeKey, error) {
	l := len(buf)
	if withTime && l < 8 {
		return nil, fmt.Errorf("wrong key buf")
	}
	name := ""
	startTime := uint64(0)
	if withTime {
		name = string(buf[0 : l-8])
		startTime = binary.BigEndian.Uint64(buf[l-8:])
	} else {
		name = string(buf)
	}
	parts := strings.Split(name, "\t")
	if len(parts) != 4 {
		return nil, fmt.Errorf("wrong key buf")
	}

	return &StartTimeKey{
		ServiceName:    parts[1],
		ServiceVersion: parts[2],
		RootSpanName:   parts[3],
		StartTime:      int64(startTime),
	}, nil
}

type _StartTimeIndex struct{}

func (index *_StartTimeIndex) Insert(txn *badger.Txn, startKey *StartTimeKey, traceId string, keepTraceDay uint) error {
	key, err := startKey.ToBytes(true)
	if err != nil {
		return err
	}
	item, err := txn.Get(key)
	if err != nil {
		if err == badger.ErrKeyNotFound {
			entry := badger.NewEntry(key, []byte(traceId)).WithTTL(time.Duration(keepTraceDay) * time.Hour * 24)
			return txn.SetEntry(entry)
		}
		return err
	}
	if item.IsDeletedOrExpired() {
		entry := badger.NewEntry(key, []byte(traceId)).WithTTL(time.Duration(keepTraceDay) * time.Hour * 24)
		return txn.SetEntry(entry)
	}

	value, err := item.ValueCopy(nil)
	if err != nil {
		return err
	}

	traceIdList := strings.Split(string(value), "\t")
	traceIdList = append(traceIdList, traceId)
	entry := badger.NewEntry(key, []byte(strings.Join(traceIdList, "\t"))).WithTTL(time.Duration(keepTraceDay) * time.Hour * 24)
	return txn.SetEntry(entry)
}

func (index *_StartTimeIndex) List(txn *badger.Txn, serviceInfo *trace_api.ServiceInfo,
	filterByRootSpanName bool, rootSpanName string,
	fromTime, toTime int64, limit uint32) ([]string, error) {
	if !filterByRootSpanName {
		rootSpanName = ""
	}
	startKey := &StartTimeKey{
		ServiceName:    serviceInfo.ServiceName,
		ServiceVersion: serviceInfo.ServiceVersion,
		RootSpanName:   rootSpanName,
	}
	prefixKey, err := startKey.ToBytes(false)
	if err != nil {
		return nil, err
	}

	options := badger.IteratorOptions{
		PrefetchValues: true,
		PrefetchSize:   100,
		Reverse:        true,
		AllVersions:    false,
		Prefix:         prefixKey,
	}
	iter := txn.NewIterator(options)
	defer iter.Close()

	iter.Rewind()
	startKey.StartTime = toTime
	key, err := startKey.ToBytes(true)
	if err != nil {
		return nil, err
	}
	iter.Seek(key)

	traceIdList := []string{}

	for ; iter.Valid(); iter.Next() {
		item := iter.Item()
		itemKey, err := ParseStartTimeKey(item.KeyCopy(nil), true)
		if err != nil {
			return nil, err
		}
		if itemKey.StartTime > toTime {
			continue
		}
		if itemKey.StartTime < fromTime {
			return traceIdList, nil
		}
		if item.IsDeletedOrExpired() {
			continue
		}
		value, err := item.ValueCopy(nil)
		if err != nil {
			return nil, err
		}
		tmpTraceIdList := strings.Split(string(value), "\t")
		traceIdList = append(traceIdList, tmpTraceIdList...)
		if len(traceIdList) >= int(limit) {
			return traceIdList, nil
		}
	}

	return traceIdList, nil
}
